import asyncio
import heapq
from typing import Dict, List, Optional, Callable, Any
from collections import defaultdict
from datetime import datetime
import logging
import uuid

from .entities import (
    SecurityLevel, 
    QueueTask, 
    CodeExecutionRequest,
    TaskStatus
)

logger = logging.getLogger(__name__)

class PriorityQueue:
    """优先级队列实现"""
    
    def __init__(self):
        self._queue = []
        self._task_map = {}  # task_id -> task
        self._counter = 0  # 用于相同优先级任务的排序
        
    def put(self, task: QueueTask[Any]) -> None:
        """添加任务到队列"""
        self._counter += 1
        # 优先级取负数，因为heapq是最小堆
        priority = -task.priority
        heapq.heappush(self._queue, (priority, self._counter, task))
        self._task_map[task.task_id] = task
        logger.debug(f"任务 {task.task_id} 已加入队列，优先级: {task.priority}")
    
    def get(self) -> Optional[QueueTask[Any]]:
        """从队列获取任务"""
        while self._queue:
            priority, counter, task = heapq.heappop(self._queue)
            if task.task_id in self._task_map:
                del self._task_map[task.task_id]
                logger.debug(f"任务 {task.task_id} 已从队列取出")
                return task
        return None
    
    def remove(self, task_id: str) -> bool:
        """从队列移除指定任务"""
        if task_id in self._task_map:
            del self._task_map[task_id]
            logger.debug(f"任务 {task_id} 已从队列移除")
            return True
        return False
    
    def size(self) -> int:
        """获取队列大小"""
        return len(self._task_map)
    
    def is_empty(self) -> bool:
        """检查队列是否为空"""
        return len(self._task_map) == 0
    
    def peek(self) -> Optional[QueueTask[Any]]:
        """查看队列头部任务但不移除"""
        while self._queue:
            priority, counter, task = self._queue[0]
            if task.task_id in self._task_map:
                return task
            else:
                # 移除已失效的任务
                heapq.heappop(self._queue)
        return None

class QueueManager:
    """队列管理器"""
    
    def __init__(self):
        # 按安全等级分别维护队列
        self.queues: Dict[SecurityLevel, PriorityQueue] = {
            SecurityLevel.LOW: PriorityQueue(),
            SecurityLevel.HIGH: PriorityQueue()
        }
        
        # 任务状态追踪
        self.task_status: Dict[str, TaskStatus] = {}
        self.task_results: Dict[str, Any] = {}
        
        # 队列统计信息
        self.stats = {
            "total_submitted": 0,
            "total_completed": 0,
            "total_failed": 0,
            "by_security_level": defaultdict(lambda: {"submitted": 0, "completed": 0, "failed": 0})
        }
        
        # 任务处理回调
        self.task_processors: Dict[SecurityLevel, Callable] = {}
        
        # 锁用于线程安全
        self._lock = asyncio.Lock()
    
    async def submit_task(self, request: CodeExecutionRequest, priority: int = 0) -> str:
        """提交任务到队列"""
        task_id = str(uuid.uuid4())
        
        task = QueueTask[CodeExecutionRequest](
            task_id=task_id,
            payload=request,
            priority=priority,
            created_at=datetime.utcnow()
        )
        
        async with self._lock:
            # 将任务添加到对应安全等级的队列
            security_level = request.security_level
            self.queues[security_level].put(task)
            
            # 更新任务状态
            self.task_status[task_id] = TaskStatus.PENDING
            
            # 更新统计信息
            self.stats["total_submitted"] += 1
            self.stats["by_security_level"][security_level]["submitted"] += 1
            
            logger.info(f"任务 {task_id} 已提交到 {security_level.value} 安全等级队列")
        
        return task_id
    
    async def get_task(self, security_level: SecurityLevel) -> Optional[QueueTask[CodeExecutionRequest]]:
        """从指定安全等级队列获取任务"""
        async with self._lock:
            queue = self.queues[security_level]
            task = queue.get()
            
            if task:
                # 更新任务状态
                self.task_status[task.task_id] = TaskStatus.RUNNING
                logger.debug(f"从 {security_level.value} 队列获取任务 {task.task_id}")
            
            return task
    
    async def complete_task(self, task_id: str, result: Any, status: TaskStatus = TaskStatus.COMPLETED) -> None:
        """标记任务完成"""
        async with self._lock:
            if task_id in self.task_status:
                old_status = self.task_status[task_id]
                self.task_status[task_id] = status
                self.task_results[task_id] = result
                
                # 更新统计信息
                if status == TaskStatus.COMPLETED:
                    self.stats["total_completed"] += 1
                elif status == TaskStatus.FAILED:
                    self.stats["total_failed"] += 1
                
                # 根据结果中的安全等级更新统计
                if hasattr(result, 'security_level'):
                    security_level = result.security_level
                    if status == TaskStatus.COMPLETED:
                        self.stats["by_security_level"][security_level]["completed"] += 1
                    elif status == TaskStatus.FAILED:
                        self.stats["by_security_level"][security_level]["failed"] += 1
                
                logger.info(f"任务 {task_id} 状态更新: {old_status} -> {status}")
    
    async def cancel_task(self, task_id: str) -> bool:
        """取消任务"""
        async with self._lock:
            # 检查任务状态
            if task_id not in self.task_status:
                return False
            
            current_status = self.task_status[task_id]
            
            # 只能取消待处理的任务
            if current_status != TaskStatus.PENDING:
                return False
            
            # 从所有队列中移除任务
            removed = False
            for security_level, queue in self.queues.items():
                if queue.remove(task_id):
                    removed = True
                    break
            
            if removed:
                self.task_status[task_id] = TaskStatus.FAILED
                self.task_results[task_id] = "任务已取消"
                logger.info(f"任务 {task_id} 已取消")
            
            return removed
    
    async def get_task_status(self, task_id: str) -> Optional[TaskStatus]:
        """获取任务状态"""
        return self.task_status.get(task_id)
    
    async def get_task_result(self, task_id: str) -> Optional[Any]:
        """获取任务结果"""
        return self.task_results.get(task_id)
    
    def get_queue_info(self, security_level: SecurityLevel) -> Dict[str, Any]:
        """获取队列信息"""
        queue = self.queues[security_level]
        return {
            "security_level": security_level.value,
            "queue_size": queue.size(),
            "is_empty": queue.is_empty(),
            "next_task": (task.task_id if (task := queue.peek()) else None)
        }
    
    def get_all_queue_info(self) -> Dict[str, Any]:
        """获取所有队列信息"""
        info = {
            "queues": {},
            "statistics": dict(self.stats)
        }
        
        for security_level in SecurityLevel:
            info["queues"][security_level.value] = self.get_queue_info(security_level)
        
        return info
    
    def register_task_processor(self, security_level: SecurityLevel, processor: Callable) -> None:
        """注册任务处理器"""
        self.task_processors[security_level] = processor
        logger.info(f"已注册 {security_level.value} 安全等级的任务处理器")
    
    async def process_tasks(self, security_level: SecurityLevel, max_concurrent: int = 5) -> None:
        """处理指定安全等级的任务"""
        if security_level not in self.task_processors:
            logger.error(f"未找到 {security_level.value} 安全等级的任务处理器")
            return
        
        processor = self.task_processors[security_level]
        semaphore = asyncio.Semaphore(max_concurrent)
        
        async def process_single_task():
            async with semaphore:
                task = await self.get_task(security_level)
                if task:
                    try:
                        result = await processor(task.payload)
                        await self.complete_task(task.task_id, result, TaskStatus.COMPLETED)
                    except Exception as e:
                        logger.error(f"处理任务 {task.task_id} 失败: {e}")
                        await self.complete_task(task.task_id, str(e), TaskStatus.FAILED)
        
        while True:
            try:
                if not self.queues[security_level].is_empty():
                    await process_single_task()
                else:
                    # 队列为空时短暂休眠
                    await asyncio.sleep(0.1)
            except Exception as e:
                logger.error(f"处理 {security_level.value} 队列任务时出错: {e}")
                await asyncio.sleep(1)
    
    async def start_queue_processors(self, max_concurrent_by_level: Dict[SecurityLevel, int] | None = None) -> List[asyncio.Task]:
        """启动所有队列处理器"""
        if max_concurrent_by_level is None:
            max_concurrent_by_level = {
                SecurityLevel.LOW: 10,
                SecurityLevel.HIGH: 3
            }
        
        tasks = []
        for security_level in SecurityLevel:
            if security_level in self.task_processors:
                max_concurrent = max_concurrent_by_level.get(security_level, 5)
                task = asyncio.create_task(
                    self.process_tasks(security_level, max_concurrent)
                )
                tasks.append(task)
                logger.info(f"已启动 {security_level.value} 队列处理器，最大并发: {max_concurrent}")
        
        return tasks
    
    def get_queue_position(self, task_id: str) -> Optional[int]:
        """获取任务在队列中的位置"""
        for security_level, queue in self.queues.items():
            if task_id in queue._task_map:
                # 简化实现：返回队列大小作为位置估算
                return queue.size()
        return None 